#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from nav_msgs.msg import Odometry
from geometry_msgs.msg import Twist, Quaternion, TransformStamped
import serial
import math
from tf2_ros import TransformBroadcaster
import re
import threading  

class Stm32OdomNode(Node):
    def __init__(self):
        super().__init__('stm32_odom_node')
        
        # 参数配置
        self.declare_parameters(
            namespace='',
            parameters=[
                ('serial_port', '/dev/ttyACM0'),
                ('baud_rate', 115200),
                ('wheel_radius', 0.1),
                ('wheel_base', 0.5),
                ('odom_frame_id', 'odom'),
                ('base_frame_id', 'base_link')
            ]
        )
        
        # ROS2发布者
        self.odom_pub = self.create_publisher(Odometry, 'odom', 10)
        self.tf_broadcaster = TransformBroadcaster(self)
        
        # 里程计状态
        self.x = 0.0
        self.y = 0.0
        self.theta = 0.0
        self.last_time = None
        
        # 初始化串口
        self.ser = serial.Serial(
            port=self.get_parameter('serial_port').value,
            baudrate=self.get_parameter('baud_rate').value,
            timeout=0.1
        )
        
        # 启动串口读取线程
        self.serial_thread = threading.Thread(target=self.read_serial)
        self.serial_thread.daemon = True
        self.serial_thread.start()
        
        self.get_logger().info("STM32里程计节点已启动")

    def read_serial(self):
        buffer = ""
        while rclpy.ok():
            try:
                data = self.ser.read(self.ser.in_waiting or 1).decode('utf-8')
                if data:
                    buffer += data
                    
                    # 处理完整数据包 (以换行符分隔)
                    while '\n' in buffer:
                        line, buffer = buffer.split('\n', 1)
                        self.process_packet(line.strip())
                        
            except Exception as e:
                self.get_logger().error(f"串口错误: {str(e)}")
                self.ser.close()
                self.ser.open()

    def process_packet(self, packet):
        """处理ODOM数据包格式: ODOM:[timestamp],L[pos],[vel],R[pos],[vel]"""
        if not packet.startswith("ODOM:"):
            return
            
        try:
            # 使用正则表达式提取数据
            match = re.match(r"ODOM:(\d+),L(-?\d+\.\d+),(-?\d+\.\d+),R(-?\d+\.\d+),(-?\d+\.\d+)", packet)
            if not match:
                return
                
            timestamp = int(match.group(1))
            left_pos = float(match.group(2))
            left_vel = float(match.group(3))
            right_pos = float(match.group(4)) 
            right_vel = -float(match.group(5))
            
            # 计算时间差 (秒)
            current_time = timestamp / 1000.0  # 假设STM32发送的是毫秒时间戳
            if self.last_time is None:
                self.last_time = current_time
                return
                
            dt = current_time - self.last_time
            self.last_time = current_time
            
            # 计算里程计 (单位: m, rad/s)
            wheel_radius = self.get_parameter('wheel_radius').value
            wheel_base = self.get_parameter('wheel_base').value
            
            v_left = left_vel * wheel_radius
            v_right = right_vel * wheel_radius
            
            # 计算机器人线速度和角速度
            vx = (v_right + v_left) / 2.0
            vtheta = (v_right - v_left) / wheel_base
            
            # 更新位姿
            self.theta += vtheta * dt
            self.x += vx * math.cos(self.theta) * dt
            self.y += vx * math.sin(self.theta) * dt
            
            # 发布里程计
            self.publish_odometry(vx, vtheta, current_time)
            
        except Exception as e:
            self.get_logger().error(f"数据解析错误: {str(e)}")

    def publish_odometry(self, vx, vtheta, timestamp):
        # 创建消息
        odom_msg = Odometry()
        tf_msg = TransformStamped()
        
        # 设置header
        odom_msg.header.stamp = self.get_clock().now().to_msg()
        odom_msg.header.frame_id = self.get_parameter('odom_frame_id').value
        odom_msg.child_frame_id = self.get_parameter('base_frame_id').value
        
        # 设置位置
        odom_msg.pose.pose.position.x = self.x
        odom_msg.pose.pose.position.y = self.y
        odom_msg.pose.pose.position.z = 0.0
        
        # 设置方向 (四元数)
        q = Quaternion()
        q.z = math.sin(self.theta / 2.0)
        q.w = math.cos(self.theta / 2.0)
        odom_msg.pose.pose.orientation = q
        
        # 设置速度
        odom_msg.twist.twist.linear.x = vx
        odom_msg.twist.twist.angular.z = vtheta
        
        # 发布消息
        self.odom_pub.publish(odom_msg)
        
        # 发布TF
        tf_msg.header = odom_msg.header
        tf_msg.child_frame_id = odom_msg.child_frame_id
        tf_msg.transform.translation.x = odom_msg.pose.pose.position.x
        tf_msg.transform.translation.y = odom_msg.pose.pose.position.y
        tf_msg.transform.translation.z = 0.0
        tf_msg.transform.rotation = odom_msg.pose.pose.orientation
        
        self.tf_broadcaster.sendTransform(tf_msg)

def main(args=None):
    rclpy.init(args=args)
    node = Stm32OdomNode()
    
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    
    node.ser.close()
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()